{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate Resnet50 Models\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import shutil\n",
    "import tensorflow as tf\n",
    "import tensorflow.neuron as tfn\n",
    "import tensorflow.compat.v1.keras as keras\n",
    "from tensorflow.keras.applications.resnet50 import ResNet50"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare export directories for compile/non-compiled versions of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dir = \"resnet50\"\n",
    "compiled_model_dir = model_dir + \"_neuron\"\n",
    "shutil.rmtree(model_dir, ignore_errors=True)\n",
    "shutil.rmtree(compiled_model_dir, ignore_errors=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instantiate a Keras ResNet50 model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/robert/.miniconda3/envs/py36-neuron/lib/python3.6/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1630: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n"
     ]
    }
   ],
   "source": [
    "keras.backend.set_learning_phase(0)\n",
    "keras.backend.set_image_data_format('channels_last')\n",
    "model = ResNet50(weights='imagenet')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Export the model as SavedModel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-4-eed95c214c0e>:5: simple_save (from tensorflow.python.saved_model.simple_save) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.simple_save.\n",
      "WARNING:tensorflow:From /home/robert/.miniconda3/envs/py36-neuron/lib/python3.6/site-packages/tensorflow_core/python/saved_model/signature_def_utils_impl.py:201: build_tensor_info (from tensorflow.python.saved_model.utils_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.utils.build_tensor_info or tf.compat.v1.saved_model.build_tensor_info.\n",
      "INFO:tensorflow:Assets added to graph.\n",
      "INFO:tensorflow:No assets to write.\n",
      "INFO:tensorflow:SavedModel written to: resnet50/saved_model.pb\n"
     ]
    }
   ],
   "source": [
    "tf.saved_model.simple_save(\n",
    "    session            = keras.backend.get_session(),\n",
    "    export_dir         = model_dir,\n",
    "    inputs             = {'input': model.inputs[0]},\n",
    "    outputs            = {'output': model.outputs[0]})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then compile it for Inferentia to be used on only one Neuron core. `--static-weights` option is used to cache all weights onto the neuron core's memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from resnet50/variables/variables\n",
      "INFO:tensorflow:Froze 320 variables.\n",
      "INFO:tensorflow:Converted 320 variables to const ops.\n",
      "INFO:tensorflow:fusing subgraph neuron_op_d6f098c01c780733 with neuron-cc\n",
      "INFO:tensorflow:Number of operations in TensorFlow session: 4638\n",
      "INFO:tensorflow:Number of operations after tf.neuron optimizations: 556\n",
      "INFO:tensorflow:Number of operations placed on Neuron runtime: 554\n",
      "INFO:tensorflow:No assets to save.\n",
      "INFO:tensorflow:No assets to write.\n",
      "INFO:tensorflow:SavedModel written to: resnet50_neuron/saved_model.pb\n",
      "INFO:tensorflow:Successfully converted resnet50 to resnet50_neuron\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'OnNeuronRatio': 0.9964028776978417}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compiler_args = ['--static-weights', '--num-neuroncores', '1']\n",
    "batch_size = 1\n",
    "tfn.saved_model.compile(model_dir, compiled_model_dir, batch_size)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
